package personal.leil.learning.prim;

import personal.leil.learning.undigraph.EdgeUndirected;
import personal.leil.learning.undigraph.Undigraph;
import personal.leil.learning.undigraph.Vertex;

import java.util.*;

/**
 * 最小生成树（Prim算法）
 * <p>
 * <pre>
 *      选择一个初始顶点作为最小生成树的根节点。
 *      重复以下步骤，直到最小生成树包含了图中的所有顶点：
 *        <ul>
 *            <li>从已经在最小生成树中的顶点集合和不在最小生成树中的顶点集合之间找到最短的边</li>
 *            <li>将该边所连接的不在最小生成树中的顶点加入最小生成树</li>
 *        </ul>
 *     </pre>
 * </p>
 *
 * @author Leil
 * @version 1.0
 * @since 2024/1/23
 */
public class PrimMST {

    public static List<EdgeUndirected> findMinimumSpanningTree(Undigraph graph) {
        Set<Vertex> set = new HashSet<>();
        List<EdgeUndirected> edges = new ArrayList<>();
        PriorityQueue<EdgeUndirected> queue = new PriorityQueue<>();
        for (Vertex vertex : graph.vertices) {
            findMinimumWeightEdgeFromVertex(vertex, set, edges, queue);
        }
        return edges;
    }


    /**
     * 找到从某一顶点出发权重最小的边（连接的顶点不在树上）
     *
     * @param vertex 目标顶点
     * @param set    已处理的集合顶点
     * @param edges  找到的目标边
     * @param queue  按权重排序的边的优先级队列
     */
    public static void findMinimumWeightEdgeFromVertex(Vertex vertex,
                                                       Set<Vertex> set,
                                                       List<EdgeUndirected> edges,
                                                       PriorityQueue<EdgeUndirected> queue) {
        if (set.contains(vertex)) return;
        for (EdgeUndirected edge : vertex.edges) {
            queue.offer(edge);
        }
        while (!queue.isEmpty()) {
            EdgeUndirected poll = queue.poll();
            Vertex another = poll.findAnother(vertex);
            if (!set.contains(another)) {
                edges.add(poll);
                set.add(another);
                findMinimumWeightEdgeFromVertex(another, set, edges, queue);
                break;
            }
        }
    }


    /************************* For Test *************************/
    public static void main(String[] args) {
        Undigraph graph = new Undigraph();
        Vertex v1 = new Vertex(1);
        graph.addVertex(v1);
        Vertex v2 = new Vertex(2);
        graph.addVertex(v2);
        Vertex v3 = new Vertex(3);
        graph.addVertex(v3);
        Vertex v4 = new Vertex(4);
        graph.addVertex(v4);

        v1.addNeighbor(v2);
        graph.addEdge(new EdgeUndirected(v1, v2, 5));
        v2.addNeighbor(v4);
        graph.addEdge(new EdgeUndirected(v2, v4, 2));
        v1.addNeighbor(v3);
        graph.addEdge(new EdgeUndirected(v1, v3, 3));
        v3.addNeighbor(v4);
        graph.addEdge(new EdgeUndirected(v3, v4, 100));
        v1.addNeighbor(v3);
        graph.addEdge(new EdgeUndirected(v1, v4, 9));
        v3.addNeighbor(v2);
        graph.addEdge(new EdgeUndirected(v3, v2, 8));

        List<EdgeUndirected> tree = findMinimumSpanningTree(graph);
        for (EdgeUndirected edge : tree) {
            System.out.println(edge);
        }

    }


}
